﻿using System;
using System.Collections.Generic;
using Algoritmia.DataStructures.Queues;
using Algoritmia.Utils;

namespace Algoritmia.DataStructures.Maps
{
    public class LeftLeaningRedBlackTreeMap<TKey,TValue> : IMap<TKey,TValue>
    {
        public class Node
        {
            public bool Red;
            public Node Left, Right;
            public TKey Key;
            public TValue Value;

            public Node(TKey key, TValue value, bool red, Node left = null, Node right = null)
            {
                Key = key;
                Value = value;
                Red = red;
                Left = left;
                Right = right;
            }

        }

        private Node _root;
        private int _count;
        private readonly Comparison<TKey> _comparison;
        private readonly Func<ILifo<Node>> _createLifo;
        private readonly KeysView<TKey, TValue> _keysView;
        private readonly ItemsView<TKey, TValue> _itemsView;
        private readonly ValuesView<TKey, TValue> _valuesView;

        public LeftLeaningRedBlackTreeMap(Comparison<TKey> comparison = null, Func<ILifo<Node>> createLifo = null)
        {
            _comparison = AlgoritmiaFactory.GetDefaultComparison(comparison);
            _createLifo = createLifo ?? (() => AlgoritmiaFactory.Lifos<Node>.Default);
            
            _keysView = new KeysView<TKey, TValue>(this);
            _itemsView = new ItemsView<TKey, TValue>(this);
            _valuesView = new ValuesView<TKey, TValue>(this);

        }

        #region Private Helper Methods

        private static bool IsRed(Node node)
        {
            return node != null && node.Red;
        }

        private Node Append(Node node,TKey key, TValue value)
        {
            if ( node == null)
            {
                _count++;
                return new Node(key, value, true);
            }

            if (IsRed(node.Left) && IsRed(node.Right))
                FlipColor(node);

            int cmp = _comparison(key, node.Key);
            if (cmp < 0) node.Left = Append(node.Left, key, value);
            if (cmp > 0) node.Right = Append(node.Right, key, value);

            if ( IsRed(node.Right) )
            {
                node = RotateLeft(node);
            }

            if ( IsRed(node.Left) && IsRed(node.Left.Left) )
            {
                node = RotateRight(node);
            }
            return node;
        }

        private static Node RotateLeft(Node node)
        {
            Node x = node.Right;
            node.Right = x.Left;
            x.Left = node;
            x.Red = node.Red;
            node.Red = true;
            return x;
        }

        private static Node RotateRight(Node node)
        { 
            Node x = node.Left;
            node.Left = x.Right;
            x.Right = node;
            x.Red = node.Red;
            node.Red = true;
            return x;
        }

        private static void FlipColor(Node node)
        {
            node.Red = !node.Red;
            node.Left.Red = !node.Left.Red;
            node.Right.Red = !node.Right.Red;
        }

        private Node Search(Node node, TKey key)
        {
            while ( node != null)
            {
                int cmp = _comparison(key, node.Key);
                if (cmp < 0) {  node = node.Left; }
                else if (cmp > 0) { node = node.Right; }
                else
                {
                    return node;
                }
            }
            return null;
        }

        private Node Delete(TKey key, Node node)
        {
            if ( _comparison(key,node.Key) < 0)
            {
                if ( node.Left != null )
                {
                    if (!IsRed(node.Left) && !IsRed(node.Left.Left))
                    {
                        node = MoveRedLeft(node);
                    }
                    node.Left = Delete(key, node.Left);
                }
            }
            else
            {
                if ( IsRed(node.Left))
                {
                    node = RotateRight(node);
                }
                if ( _comparison(key,node.Key) == 0 && node.Right == null)
                {
                    _count--;
                    return null;
                }
                if (node.Right != null)
                {
                    if (!IsRed(node.Right) && !IsRed(node.Right.Left))
                    {
                        node = MoveRedRight(node);
                    }
                    if ( _comparison(key,node.Key) == 0)
                    {
                        Node minimum = SearchMin(node.Right);
                        node.Key = minimum.Key;
                        node.Value = minimum.Value;
                        node.Right = DeleteMin(node.Right);
                    }
                    else
                    {
                        node.Right = Delete(key, node.Right);
                    }
                }

            }
            return FixUp(node);
        }

        private static Node FixUp(Node node)
        {
            if (IsRed(node.Right))
            {
                node = RotateLeft(node);
            }
            if (IsRed(node.Left) && IsRed(node.Left.Left))
            {
                node = RotateRight(node);
            }
            if (IsRed(node.Left) && IsRed(node.Right))
            {
                FlipColor(node);
            }
            if ( node.Left != null && IsRed(node.Left.Right) && !IsRed(node.Left.Left))
            {
                node.Left = RotateLeft(node.Left);
                if ( IsRed(node.Left))
                {
                    node = RotateRight(node);
                }
            }
            return node;
        }
        
        private Node DeleteMin(Node node)
        {
            if (node == null)
                return null;
            if (!IsRed(node.Left) && !IsRed(node.Left.Left))
            {
                node = MoveRedLeft(node);
            }
            node.Left = DeleteMin(node.Left);
            _count++;
            return FixUp(node);
        }

        private static Node SearchMin(Node node)
        {
            while (node.Left != null)
            {
                node = node.Left;
            }
            return node;
        }

        private static Node MoveRedLeft(Node node)
        {
            FlipColor(node);
            if (IsRed((node.Right.Left)))
            {
                node.Right = RotateRight(node.Right);
                node = RotateLeft(node);
                FlipColor(node);
                if (IsRed(node.Right.Right))
                {
                    node.Right = RotateLeft(node.Right);
                }
            }
            return node;
        }

        private static Node MoveRedRight(Node node)
        {
            FlipColor(node);
            if ( IsRed(node.Left.Left))
            {
                node = RotateRight(node);
                FlipColor(node);
            }
            return node;
        }

        private IEnumerable<Node> InorderTraversal(Node node)
        {
            ILifo<Node> stack = _createLifo();
            Node current = node;
            while ( current != null )
            {
                if ( current.Left != null)
                {
                    stack.Push(current);
                    current = current.Left;
                }
                else
                {
                    while ( current != null)
                    {
                        yield return current;
                        current = current.Right;
                        if (current != null) break;
                        if (stack.Count == 0) break;
                        current = stack.Pop();
                    }
                }
            }
        }

        #endregion

        #region IMap<TKey,TValue> Members

        public TValue this[TKey key]
        {
            get
            {
                Node node = Search(_root, key);
                if ( node == null)
                {
                    throw new KeyNotFoundException(String.Format("Not found any key equals than {0}", key));
                }
                return node.Value; 
            }
            set
            {
                _root = Append(_root, key, value);
                _root.Red = false;
            }
        }

        public void Add(TKey key, TValue value)
        {
            this[key] = value;
        }

        public bool Contains(TKey key)
        {
            return Search(_root, key) != null;
        }

        public TValue Get(TKey key, TValue defaultValue)
        {
            Node node = Search(_root, key);
            if ( node == null)
                return defaultValue;
            return node.Value;
        }

        public ICollectionView<TKey> Keys()
        {
            return _keysView;
        }

        public ICollectionView<KeyValuePair<TKey, TValue>> Items()
        {
            return _itemsView;
        }

        public bool Remove(TKey key)
        {
            int prev = _count-1;
            _root = Delete(key,_root);
            if (_root != null) _root.Red = false;
            return prev == _count;
        }

       
        public TValue SetDefault(TKey key, TValue defaultValue)
        {
            Node node = Search(_root, key);
            if (node == null)
            {
                this[key] = defaultValue;
                return defaultValue;
            }
            return node.Value;
        }

        public ICollectionView<TValue> Values()
        {
            return _valuesView;
        }

        #endregion

        #region IEnumerable<KeyValuePair<TKey,TValue>> Members

        public IEnumerator<KeyValuePair<TKey, TValue>> GetEnumerator()
        {
            foreach(var node in InorderTraversal(_root))
            {
                yield return new KeyValuePair<TKey, TValue>(node.Key,node.Value);
            }
        }

        #endregion

        #region IEnumerable Members

        System.Collections.IEnumerator System.Collections.IEnumerable.GetEnumerator()
        {
            return GetEnumerator();
        }

        #endregion

        #region ICountable Members

        public int Count
        {
            get { return _count; }
        }

        #endregion
    }
}
